#!/usr/bin/env python3

from __future__ import annotations

import argparse
import logging
import os
import resource
import signal
import sys
import time

import logger
from model import Settings, Action, Request


DESCRIPTION = "Tool for ChinaSoft 2025 Heterogeneous Memory Challenge."
CPU_LIMIT = int(os.environ.get("CPU_LIMIT", 10))  # 10 seconds
RAM_LIMIT = int(os.environ.get("RAM_LIMIT", 1)) * 1024**3  # 1 GiB


def file_path(parser, s, allow_stdin: bool = True) -> str:
    if os.path.isfile(s) or (allow_stdin and s == "-"):
        return s
    parser.error(f"File '{s}' does not exist.")


def parse_extra_args(parser, extra_args):
    kwargs = {}
    key = None

    for arg in extra_args:
        if arg.startswith("-"):
            key = arg.lstrip("-")
            kwargs[key] = True
        else:
            if key is None:
                parser.error(f"Unexpected argument '{arg}' without a preceding key.")
            kwargs[key] = arg
            key = None

    return kwargs


def set_resource_limits(memory_limit=RAM_LIMIT):
    """
    Set resource limits for memory usage.

    Args:
        memory_limit: Maximum memory in bytes (default: RAM_LIMIT)
    """
    try:
        # Set memory limit (in bytes)
        resource.setrlimit(resource.RLIMIT_AS, (memory_limit, memory_limit))
    except (ValueError, OSError) as e:
        logger.logger.warning(f"Failed to set resource limits: {e}")


class TimeoutException(Exception):
    """Exception raised when algorithm execution times out."""

    pass


def timeout_handler(signum, frame):
    """Handler for SIGALRM - raises TimeoutException."""
    raise TimeoutException("Algorithm exceeded time limit")


def common_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=DESCRIPTION)
    parser.add_argument(
        "-c",
        "--color",
        choices=["auto", "always", "never"],
        help="Colorize the output. (default: auto)",
        default="auto",
    )
    parser.add_argument(
        "-i",
        "--input",
        required=True,
        type=lambda s: file_path(parser, s),
        help="Input file for the tool. (or '-' for stdin)",
    )
    return parser


def do_data_gen() -> int:
    from generator import RequestBuilder, TickGenerator

    parser = argparse.ArgumentParser(description=DESCRIPTION)
    parser.add_argument(
        "-c",
        "--color",
        choices=["auto", "always", "never"],
        help="Colorize the output. (default: auto)",
        default="auto",
    )
    parser.add_argument(
        "-i",
        "--input",
        required=True,
        help="Input file for the tool. (or '-' for stdout)",
    )
    parser.add_argument(
        "-f",
        "--force",
        action="store_true",
        help="Overwrite the input file if it already exists.",
    )
    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
    parser.add_argument(
        "-m",
        "--mode",
        choices=TickGenerator.list_modes(),
        default=TickGenerator.list_modes()[0],
        help="Mode of data generation.",
    )
    parser.add_argument(
        "-g",
        "--max-time-gap",
        type=int,
        default=100,
        help="Max time gap between two consecutive requests.",
    )
    parser.add_argument(
        "-L",
        "--virtual-memory-size",
        type=int,
        default=1000,
        help="Virtual memory size.",
    )
    parser.add_argument(
        "-M",
        "--physical-memory-size",
        type=int,
        default=800,
        help="Physical memory size.",
    )
    parser.add_argument(
        "-N",
        "--num-requests",
        type=int,
        default=10,
        help="Number of requests to generate.",
    )
    parser.add_argument("extra_args", nargs=argparse.REMAINDER)

    args = parser.parse_args()
    extra_args = parse_extra_args(parser, args.extra_args)

    if os.path.exists(args.input) and args.input != "-":
        if args.force:
            import shutil

            if os.path.isfile(args.input):
                os.remove(args.input)
            else:
                shutil.rmtree(args.input)
        else:
            parser.error(f"Path '{args.input}' already exists.")

    logger.init(args.color)

    settings = Settings(
        args.virtual_memory_size, args.physical_memory_size, args.num_requests
    )
    generator = TickGenerator.from_mode(args.mode, settings, **extra_args)
    builder = RequestBuilder(settings)
    requests = builder.build(generator, args.max_time_gap, seed=args.seed)
    settings.N = len(requests)

    with (
        sys.stdout if args.input == "-" else open(args.input, "w", encoding="utf-8")
    ) as input_textio:
        Request.into_textio(input_textio, settings, requests)

    return 0


def do_migrate() -> int:
    from migrator import Migrator

    parser = argparse.ArgumentParser(description=DESCRIPTION)
    parser.add_argument("inputs", nargs="+", help="Input files to migrate.")
    parser.add_argument(
        "-m",
        "--migration",
        choices=Migrator.list_types(),
        required=True,
        help="Migration type.",
    )
    args = parser.parse_args()

    migrator = Migrator.from_type(args.migration)
    for inp in args.inputs:
        with open(inp, "r", encoding="utf-8") as input_textio:
            settings, requests = Request.from_textio(input_textio)
        migrated_requests = migrator.migrate(settings, requests)
        with open(inp, "w", encoding="utf-8") as output_textio:
            Request.into_textio(output_textio, settings, migrated_requests)

    return 0


def do_solve() -> int:
    from algorithms import Algorithm

    parser = common_parser()
    parser.add_argument(
        "-o",
        "--output",
        required=True,
        help="Output file for the solution. (or '-' for stdout)",
    )
    parser.add_argument(
        "-f",
        "--force",
        action="store_true",
        help="Overwrite the output file if it already exists.",
    )
    parser.add_argument(
        "-a",
        "--algorithm",
        choices=Algorithm.list_names(),
        default=Algorithm.best_name(),
        help="Algorithm to use for solving the problem.",
    )
    parser.add_argument(
        "--check-limit",
        action="store_true",
        help="Check resource limits during execution.",
    )
    args = parser.parse_args()

    if args.input == "-" and args.output == "-":
        parser.error("Input and output cannot both be stdin ('-').")

    if os.path.exists(args.output) and args.output != "-":
        if args.force:
            import shutil

            if os.path.isfile(args.output):
                os.remove(args.output)
            else:
                shutil.rmtree(args.output)
        else:
            parser.error(f"Path '{args.output}' already exists.")

    logger.init(args.color)

    def run():
        now = time.time()
        with (
            sys.stdin if args.input == "-" else open(args.input, "r", encoding="utf-8")
        ) as input_textio:
            settings, requests = Request.from_textio(input_textio)

        algorithm = Algorithm.from_name(args.algorithm, settings)

        actions = algorithm.fit(requests)
        with (
            sys.stdout
            if args.output == "-"
            else open(args.output, "w", encoding="utf-8")
        ) as output_textio:
            Action.into_textio(output_textio, actions)
        elapsed = time.time() - now
        logger.logger.info(f"Algorithm finished in {elapsed:.2f} seconds 🚀")

    if args.check_limit:
        limit_state = None
        set_resource_limits(memory_limit=RAM_LIMIT)
        try:
            old_handler = signal.signal(signal.SIGALRM, timeout_handler)
            signal.alarm(CPU_LIMIT)

            try:
                run()
            finally:
                # Cancel the alarm and restore old handler
                signal.alarm(0)
                signal.signal(signal.SIGALRM, old_handler)

        except TimeoutException:
            limit_state = "CPU"
            return 1
        except MemoryError:
            limit_state = "Memory"
            return 1
        finally:
            # Log the state here instead of in the except blocks to avoid memory limit issues
            if limit_state == "CPU":
                logger.logger.error("❌ Algorithm exceeded CPU limit")
            elif limit_state == "Memory":
                logger.logger.error("❌ Algorithm exceeded memory limit")
    else:
        run()

    return 0


def do_validate() -> int:
    from validator import Validator

    parser = common_parser()
    parser.add_argument(
        "-o",
        "--output",
        type=lambda s: file_path(parser, s),
        help="Output file for validation. If not provided, only input validation is performed. \
(or '-' for stdin)",
    )
    parser.add_argument(
        "--current-best",
        type=str,
        help="File containing the current best solution to compare against.",
    )
    parser.add_argument(
        "--strict-best",
        action="store_true",
        help="Fail if not better than current best.",
    )
    parser.add_argument(
        "-v", "--verbose", action="store_true", help="Enable verbose output."
    )
    args = parser.parse_args()

    if args.input == "-" and args.output == "-":
        parser.error("Input and output cannot both be stdin ('-').")

    logger.init(args.color)

    if args.verbose:
        logger.logger.setLevel(logging.DEBUG)

    with (
        sys.stdin if args.input == "-" else open(args.input, "r", encoding="utf-8")
    ) as input_textio:
        settings, requests = Request.from_textio(input_textio)
    if args.output:
        with (
            sys.stdin
            if args.output == "-"
            else open(args.output, "r", encoding="utf-8")
        ) as output_textio:
            actions = Action.from_textio(settings, requests, output_textio)
    else:
        actions = None
    validator = Validator(settings, requests, actions, verbose=args.verbose)
    if not validator.validate():
        return 1

    if args.current_best and actions:
        fin_time = actions[-1].time
        if os.path.exists(args.current_best):
            with open(args.current_best, "r", encoding="utf-8") as best_textio:
                best_time = int(best_textio.readline().strip())
            if fin_time > best_time:
                if args.strict_best:
                    logger.logger.error(
                        f"Solution is not better than the current best. ❌ {fin_time} > {best_time}"
                    )
                    return 1
                else:
                    logger.logger.warning(
                        f"Solution is not better than the current best. ⚠️ {fin_time} > {best_time}"
                    )
            elif fin_time < best_time:
                logger.logger.info(f"New best solution found! 🎉 {fin_time} < {best_time}")
                with open(args.current_best, "w", encoding="utf-8") as best_textio:
                    best_textio.write(f"{actions[-1].time}\n")
        else:
            with open(args.current_best, "w", encoding="utf-8") as best_textio:
                best_textio.write(f"{actions[-1].time}\n")
    return 0


def do_visualise() -> int:
    from validator import Validator
    from visualizer import Visualizer

    parser = common_parser()
    parser.add_argument(
        "-o",
        "--output",
        required=True,
        type=lambda s: file_path(parser, s),
        help="Output file for validation. (or '-' for stdin)",
    )
    parser.add_argument(
        "-p",
        "--path",
        type=str,
        required=True,
        help="Path to save the visualization output. Suggest an SVG file.",
    )
    parser.add_argument(
        "-f",
        "--force",
        action="store_true",
        help="Overwrite the output path if it already exists.",
    )
    parser.add_argument(
        "--trust",
        action="store_true",
        help="Trust that the input and output are valid. Skips validation step.",
    )
    args = parser.parse_args()

    if args.input == "-" and args.output == "-":
        parser.error("Input and output cannot both be stdin ('-').")

    if os.path.exists(args.path):
        if args.force:
            import shutil

            if os.path.isfile(args.path):
                os.remove(args.path)
            else:
                shutil.rmtree(args.path)
        else:
            parser.error(f"Path '{args.path}' already exists.")

    logger.init(args.color)

    with (
        sys.stdin if args.input == "-" else open(args.input, "r", encoding="utf-8")
    ) as input_textio:
        settings, requests = Request.from_textio(input_textio)
    with (
        sys.stdin if args.output == "-" else open(args.output, "r", encoding="utf-8")
    ) as output_textio:
        actions = Action.from_textio(settings, requests, output_textio)

    if not args.trust:
        validator = Validator(settings, requests, actions, verbose=False)
        if not validator.validate():
            return 1

    visualizer = Visualizer(settings, requests, actions)
    visualizer.visualize(args.path)
    return 0


def do_compare() -> int:
    import re

    from comparator import Comparator

    parser = argparse.ArgumentParser(description=DESCRIPTION)
    parser.add_argument(
        "-i",
        "--input",
        required=True,
        type=lambda s: file_path(parser, s, allow_stdin=False),
        help="Input file for the tool.",
    )
    parser.add_argument(
        "-o",
        "--outputs",
        required=True,
        nargs="+",
        type=lambda s: file_path(parser, s, allow_stdin=False),
        help="Output files to compare.",
    )
    parser.add_argument(
        "--input-pattern",
        type=str,
        default="{name}.in",
        help="Input file pattern with '{name}' placeholder.",
    )
    parser.add_argument(
        "--output-pattern",
        type=str,
        default="{name}.{algo}.out",
        help="Output file pattern with '{name}' and '{algo}' placeholders.",
    )
    parser.add_argument(
        "-s",
        "--style",
        type=str,
        default="markdown",
        choices=Comparator.list_styles(),
        help="Output style for comparison table.",
    )
    args = parser.parse_args()

    input_re = re.escape(args.input_pattern).replace(r"\{name\}", r"(?P<name>.+)")
    output_re = (
        re.escape(args.output_pattern)
        .replace(r"\{name\}", r"(?P<name>.+)")
        .replace(r"\{algo\}", r"(?P<algo>.+)")
    )
    input_match = re.match(input_re, args.input)
    if not input_match:
        parser.error(
            f"Input file '{args.input}' does not match the input pattern '{args.input_pattern}'."
        )
    input_name = input_match.group("name")
    with open(args.input, "r", encoding="utf-8") as input_textio:
        settings, requests = Request.from_textio(input_textio)

    actions_dict = {}
    for output_file in args.outputs:
        output_match = re.match(output_re, output_file)
        if not output_match:
            parser.error(
                f"Output file '{output_file}' does not match the output pattern '{args.output_pattern}'."
            )
        output_name = output_match.group("name")
        algo_name = output_match.group("algo")
        if output_name != input_name:
            parser.error(
                f"Output file '{output_file}' name '{output_name}' does not match input name '{input_name}'."
            )
        with open(output_file, "r", encoding="utf-8") as output_textio:
            actions = Action.from_textio(settings, requests, output_textio)
        actions_dict[algo_name] = actions

    comparator = Comparator(settings, requests, actions_dict, args.style)
    sys.stdout.write(comparator.emit_table())
    return 0


if __name__ == "__main__":
    match os.path.basename(sys.argv[0]):
        case "data-gen":
            exit(do_data_gen())
        case "migrate":
            exit(do_migrate())
        case "solve":
            exit(do_solve())
        case "validate":
            exit(do_validate())
        case "visualise":
            exit(do_visualise())
        case "compare":
            exit(do_compare())
        case _:
            print(
                "Please invoke one of the specific tools: data-gen, migrate, solve, validate, visualise, compare.",
                file=sys.stderr,
            )
            exit(1)
